import datetime
import logging

import graphviz
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import accuracy_score, confusion_matrix, ConfusionMatrixDisplay
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.preprocessing import OneHotEncoder
from sklearn.tree import DecisionTreeClassifier, export_graphviz

# 配置日志
logging.basicConfig(filename=f'./logs/decision_tree.log',
                    level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    encoding='utf-8')

# 加载数据集
# dataset_path = './data/car_1000.txt'
dataset_path = './data/car_1000_numeric.csv'
logging.info(f"正在加载数据集: {dataset_path}")
data = pd.read_csv(dataset_path, header=None)

# 分离特征和标签
X = data.iloc[:, :-1]  # 特征
y = data.iloc[:, -1]  # 标签

# 处理类别变量（OneHot编码）
# encoder = OneHotEncoder()
# X_encoded = encoder.fit_transform(X)

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建决策树分类器
clf = DecisionTreeClassifier()

# 实现后剪枝
clf_post_pruned = DecisionTreeClassifier()
path = clf_post_pruned.cost_complexity_pruning_path(X_train, y_train)
ccp_alphas, impurities = path.ccp_alphas, path.impurities

fig, ax = plt.subplots(figsize=(16,8))
ax.plot(ccp_alphas[:-1], impurities[:-1], marker='o', drawstyle="steps-post")
# fig.savefig('./decision_tree/decision_tree_cost_complexity_pruning_path.png')


clfs = []
for ccp_alpha in ccp_alphas:
    clf = DecisionTreeClassifier(ccp_alpha=ccp_alpha)
    clf.fit(X_train, y_train)
    clfs.append(clf)
scores = [clf.score(X_test, y_test) for clf in clfs]
best_alpha = ccp_alphas[scores.index(max(scores))]


clf_post_pruned = DecisionTreeClassifier(ccp_alpha=best_alpha)
clf_post_pruned.fit(X_train, y_train)
y_pred_post_pruned = clf_post_pruned.predict(X_test)
accuracy_post_pruned = accuracy_score(y_test, y_pred_post_pruned)

print(f"参数：{clf_post_pruned.get_params()}")
logging.info(f"参数：{clf_post_pruned.get_params()}")
print("后剪枝准确率:", accuracy_post_pruned)
logging.info(f"后剪枝准确率: {accuracy_post_pruned}")

# 可视化 https://zhuanlan.zhihu.com/p/268532582
dot_data = export_graphviz(clf_post_pruned, filled=True,
                           # 这里还没调通
                           # feature_names=['buying', 'maint', 'doors', 'persons', 'lug_boot', 'safety'],
                           class_names=['unacc', 'acc', 'good', 'vgood'],
                           label='root', node_ids=True)
graph = graphviz.Source(dot_data, format='png')
graph.render(directory="./decision_tree/",
             filename=f'{datetime.datetime.now().strftime("%Y-%m-%d_%H_%M_%S")}后剪枝决策树可视化')

# 绘制混淆矩阵
conf_matrix = confusion_matrix(y_test, y_pred_post_pruned)
ConfusionMatrixDisplay.from_estimator(clf_post_pruned, X_test, y_test, cmap=plt.get_cmap('Blues'))
plt.title('Post pruned confusion matrix')
plt.savefig(f'./decision_tree/{datetime.datetime.now().strftime("%Y-%m-%d_%H_%M_%S")}后剪枝混淆矩阵.png')
plt.show()
